/*
 * This file is part of JSTUN. 
 * 
 * Copyright (c) 2005 Thomas King <king@t-king.de> - All rights
 * reserved.
 * 
 * This software is licensed under either the GNU Public License (GPL),
 * or the Apache 2.0 license. Copies of both license agreements are
 * included in this distribution.
 */
package com.peerialism.natcracker.discovery;

import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.net.UnknownHostException;

import org.apache.log4j.Logger;

import static com.peerialism.natcracker.discovery.DiscoveryConstants.ALLOWED_CONTIGUITY_INCREMENT;
import static com.peerialism.natcracker.discovery.DiscoveryConstants.MAX_TIMEOUT;
import static com.peerialism.natcracker.discovery.DiscoveryConstants.TIME_SINCE_FIRST_TRANSMISSION;

import com.peerialism.natcracker.common.ClashingPolicy;
import com.peerialism.natcracker.common.FilteringPolicy;
import com.peerialism.natcracker.common.PortAllocationPolicy;
import com.peerialism.natcracker.common.PortMappingPolicy;
import com.peerialism.natcracker.discovery.jstun.attribute.ChangeRequest;
import com.peerialism.natcracker.discovery.jstun.attribute.ChangeResponse;
import com.peerialism.natcracker.discovery.jstun.attribute.ChangedAddress;
import com.peerialism.natcracker.discovery.jstun.attribute.ErrorCode;
import com.peerialism.natcracker.discovery.jstun.attribute.MappedAddress;
import com.peerialism.natcracker.discovery.jstun.attribute.MessageAttribute;
import com.peerialism.natcracker.discovery.jstun.attribute.MessageAttributeException;
import com.peerialism.natcracker.discovery.jstun.attribute.MessageAttributeParsingException;
import com.peerialism.natcracker.discovery.jstun.header.MessageHeader;
import com.peerialism.natcracker.discovery.jstun.header.MessageHeaderParsingException;
import com.peerialism.natcracker.discovery.jstun.utils.UtilityException;

public class DiscoveryTest {
	/*
	 * All STUN servers have response time < 60 ms, this should then be a good value
	 */
	private static Logger logger = Logger.getLogger(DiscoveryTest.class);
	static int TIMEOUT_INIT_VALUE = 300; // ms
	InetAddress iaddress;
	int originalPort;
	private InetAddress originalAddress;

	public DiscoveryTest(InetAddress iaddress, String stunServer) {
		this.iaddress = iaddress;
		int index = stunServer.indexOf(":");
		String stunServer1 = stunServer.substring(0, index);
		originalPort = Integer.parseInt(stunServer.substring(index + 1, stunServer.length()));
		try {
			this.originalAddress = InetAddress.getByName(stunServer1);
		} catch (UnknownHostException e) {
			logger.warn("Cannot resolve " + stunServer1);
		}
	}

	public DiscoveryInfoBehave test(DatagramSocket sock) throws UtilityException, SocketException, UnknownHostException, IOException, MessageAttributeParsingException, MessageAttributeException,
			MessageHeaderParsingException {
		DiscoveryInfoBehave di = new DiscoveryInfoBehave(iaddress);
		// Use the local port specified in the call, for UPnP (-1 if random)
		di.setLocalPortUsed(sock.getLocalPort());
		InetSocketAddress changedAddress = null;
		if ((changedAddress = testPortAllocation(sock, di)) != null) {
			testFiltering(sock, di, changedAddress);
			testMappingTechnique(sock, di);
			if (di.getNatDevice().getPortAllocationPolicy() == PortAllocationPolicy.PORT_PRESERVATION && di.getNatDevice().getPortMappingPolicy() == PortMappingPolicy.PORT_DEPENDENT) {
				testClashing(sock, di);
			}
			if (di.getNatDevice().getPortFilteringPolicy() == FilteringPolicy.ENDPOINT_INDEPENDENT && di.getNatDevice().getPortMappingPolicy() == PortMappingPolicy.ENDPOINT_INDEPENDENT) {
				testSHPMultisession(sock, di, changedAddress);
			}
		}
		return di;
	}

	private InetSocketAddress testPortAllocation(DatagramSocket sock, DiscoveryInfoBehave di) throws SocketTimeoutException, SocketException, MessageAttributeParsingException, MessageHeaderParsingException,
			UtilityException, IOException, MessageAttributeException {
		traceLog("Port Mapping Test");
		traceLog("Test1");
		// Test I
		ChangeResponse response1 = changeRequest(sock, di, iaddress, di.getLocalPortUsed(), originalAddress, originalPort);
		// If Test I succeeded
		if (!response1.isNoResponse()) {
			di.setPublicIP(response1.getMappedAddress().getAddress());
			// If Mapped address == local address => OpenInternet
			if (response1.getMappedAddress().getAddress().equals(iaddress) && response1.getMappedPort() == response1.localPort) {
				/* Open */
				di.setOpenAccess();
				traceLog("OpenInternet");
			} else {
				/* Natted */
				di.setNatted(true);
				di.setLocalPortUsed(response1.localPort);
				traceLog("Natted");
				traceLog("Test2");
				// Test II - Send to primary port changed address
				ChangeResponse response2 = changeRequest(sock, di, iaddress, response1.localPort, response1.getChangedAddress().getAddress(), originalPort);
				if (!response2.isNoResponse()) {
					if (response2.getMappedAddress().equals(response1.getMappedAddress())) {
						// Endpoint independent
						// di.getNatDevice().setPortMappingPolicy(PortMappingPolicy.ENDPOINT_INDEPENDENT);
						// Test II.a - Send to primary port changed address
						traceLog("Test2.A");
						ChangeResponse response2a = changeRequest(sock, di, iaddress, response1.localPort, response1.getChangedAddress().getAddress(), response1.getChangedPort());
						if (response2.getMappedAddress().equals(response2a.getMappedAddress())) {
							// Endpoint independent
							di.getNatDevice().setPortMappingPolicy(PortMappingPolicy.ENDPOINT_INDEPENDENT);
							traceLog("Endpoint Independent");
						} else {
							// Port dependent
							di.getNatDevice().setPortMappingPolicy(PortMappingPolicy.PORT_DEPENDENT);
							traceLog("Port Dependent");
						}
					} else {
						// Test III
						traceLog("Test3");
						ChangeResponse response3 = changeRequest(sock, di, iaddress, response1.localPort, response1.getChangedAddress().getAddress(), response1.getChangedPort());
						if (!response3.isNoResponse()) {
							if (response3.getMappedAddress().equals(response2.getMappedAddress())) {
								// Endpoint dependent
								di.getNatDevice().setPortMappingPolicy(PortMappingPolicy.HOST_DEPENDENT);
								traceLog("Endpoint Dependent");
							} else {
								// Port dependent
								di.getNatDevice().setPortMappingPolicy(PortMappingPolicy.PORT_DEPENDENT);
								traceLog("Port Dependent");
							}
						} else {
							// Test III failed
							traceLog("Port Mapping: Test III failed");
							di.setBlockedUDP();
						}
					}
				} else {
					// Test II failed
					traceLog("Port Mapping: Test II failed");
					di.setBlockedUDP();
				}
			}
		} else {
			traceLog("Port Mapping: Test I failed");
			di.setBlockedUDP();
		}
		return response1.changedAddress;
	}

	private void testFiltering(DatagramSocket sock, DiscoveryInfoBehave di, InetSocketAddress changedAddress) throws SocketException, MessageAttributeParsingException, MessageHeaderParsingException, UtilityException,
			IOException, MessageAttributeException {
		traceLog("Port Filtering Test");
		traceLog("Test 2");
		traceLog("Send to: " + originalAddress + ":" + originalPort + " expect reply from: " + changedAddress);
		// Test II -- Set change address and port
		ChangeResponse response1 = changeRequest(sock, di, iaddress, di.getLocalPortUsed(), originalAddress, originalPort, false, true, changedAddress.getAddress(), changedAddress.getPort());
		if (!response1.isNoResponse()) {
			// Endpoint Independent
			di.getNatDevice().setPortFilteringPolicy(FilteringPolicy.ENDPOINT_INDEPENDENT);
			traceLog("Endpoint Independent");
		} else {
			if (di.isOpenAccess()) {
				di.setNatted(true);
			}
			traceLog("Test 3");
			traceLog("Send to: " + originalAddress + ":" + originalPort + " expect reply from: " + changedAddress.getAddress() + ":" + originalPort);
			// Test III -- Set change port, same IP
			ChangeResponse response2 = changeRequest(sock, di, iaddress, response1.localPort, originalAddress, originalPort, false, true, changedAddress.getAddress(), changedAddress.getPort());
			if (response2.isNoResponse()) {
				di.getNatDevice().setPortFilteringPolicy(FilteringPolicy.PORT_DEPENDENT);
				traceLog("Address Port Dependent");
			} else {
				di.getNatDevice().setPortFilteringPolicy(FilteringPolicy.HOST_DEPENDENT);
				traceLog("Address Dependent");
			}
		}
	}

	private void testMappingTechnique(DatagramSocket sock, DiscoveryInfoBehave di) throws UtilityException, SocketException, UnknownHostException, IOException, MessageAttributeParsingException, MessageAttributeException,
			MessageHeaderParsingException {
		// Try first IP
		traceLog("Port Prediction Test");
		ChangeResponse response = changeRequest(sock, di, iaddress, di.getLocalPortUsed(), originalAddress, originalPort);
		if (!response.isNoResponse()) {
			int publicPort1 = response.getMappedPort();
			int localPort1 = response.localPort;
			// Try second IP
			response = changeRequest(sock, di, iaddress, -1, response.getChangedAddress().getAddress(), response.getChangedPort());
			int publicPort2 = response.getMappedPort();
			int localPort2 = response.localPort;
			if (localPort2 == publicPort2 && localPort1 == publicPort1) {
				// If two ports match -- port preservation
				di.getNatDevice().setPortAllocationPolicy(PortAllocationPolicy.PORT_PRESERVATION);
				traceLog("Test Behave: Port Preservation supported, local port=" + localPort2 + ", public port=" + publicPort2);
			} else {
				traceLog("Test Behave: Port Preservation not supported, local port=" + localPort2 + ", public port=" + publicPort2);
				// Do check for correlation in port contiguity
				int p = publicPort2 - publicPort1;
				if (0 < p && p <= ALLOWED_CONTIGUITY_INCREMENT) {
					di.getNatDevice().setPortAllocationPolicy(PortAllocationPolicy.PORT_CONTIGUITY);
					di.getNatDevice().setPortIncrement(p);
					traceLog("Test Behave: Port Contiguity supported, public port1=" + publicPort1 + ", public port2=" + publicPort2);
				} else {
					di.getNatDevice().setPortAllocationPolicy(PortAllocationPolicy.RANDOM);
					traceLog("Test Behave: Port Contiguity not supported, public port1=" + publicPort1 + ", public port2=" + publicPort2);
				}
			}
		} else {
			traceLog("Test Port preservation failed ");
			return;
		}
	}

	/* Policy-specific tests */
	private void testSHPMultisession(DatagramSocket sock, DiscoveryInfoBehave di, InetSocketAddress changedAddress) throws SocketException, MessageAttributeParsingException, MessageHeaderParsingException,
			UtilityException, IOException, MessageAttributeException {
		traceLog("SHP Multisession Test");
		traceLog("Send to: " + originalAddress + ":" + originalPort + " expect reply from: " + changedAddress);
		// Test II -- Set change address and port
		for (int i = 0; i < DiscoveryConstants.SHP_MULTISESSION_ITERATIONS; i++) {
			ChangeResponse response1 = changeRequest(sock, di, iaddress, di.getLocalPortUsed(), originalAddress, originalPort, false, false, changedAddress.getAddress(), changedAddress.getPort());
			if (response1.isNoResponse()) {
				return;
			}
			ChangeResponse response2 = changeRequest(sock, di, iaddress, response1.localPort, changedAddress.getAddress(), changedAddress.getPort(), false, false, changedAddress.getAddress(),
					changedAddress.getPort());
			if (response2.isNoResponse()) {
				return;
			}
		}
		traceLog("SHP Multisession Supported");
		di.getNatDevice().setSHPMultisessionSupported(true);
	}

	private void testClashing(DatagramSocket sock, DiscoveryInfoBehave di) throws UtilityException, SocketException, UnknownHostException, IOException, MessageAttributeParsingException, MessageAttributeException,
			MessageHeaderParsingException {
		traceLog("Clashing Test");
		ChangeResponse response = changeRequest(sock, di, iaddress, -1, originalAddress, originalPort);
		if (!response.isNoResponse()) {
			ChangeResponse responseSecond = changeRequest(sock, di, iaddress, -1, originalAddress, response.getChangedPort());
			if (!responseSecond.isNoResponse()) {
				if (response.getMappedPort() != responseSecond.getMappedPort()) {
					di.getNatDevice().setPortClashingPolicy(ClashingPolicy.OPENS);
				} else {
					di.getNatDevice().setPortClashingPolicy(ClashingPolicy.OVERRIDES);
				}
			}
		}
	}

	public static ChangeResponse changeRequest(DatagramSocket sock, DiscoveryInfoBehave di, InetAddress localAddress, int localPort, InetAddress stunServer, int port) throws SocketTimeoutException, SocketException,
			UnknownHostException, MessageAttributeParsingException, MessageHeaderParsingException, UtilityException, IOException, MessageAttributeException {
		return changeRequest(sock, di, localAddress, localPort, stunServer, port, false, false, null, -1);
	}

	private static ChangeResponse changeRequest(DatagramSocket sock, DiscoveryInfoBehave di, InetAddress localAddress, int localPort, InetAddress stunServer, int port, boolean changeAddress, boolean changePort,
			InetAddress changedAddress, int changedPort) throws SocketTimeoutException, UtilityException, SocketException, UnknownHostException, IOException, MessageAttributeParsingException,
			MessageAttributeException, MessageHeaderParsingException {
		int timeout = TIMEOUT_INIT_VALUE;
		int timeSinceFirstTransmission = 0;
		//InetAddress remoteAddressTemp = stunServer;
		//int remotePortTemp = port;
		traceLog(localAddress, "Contacting STUN server: " + stunServer + ":" + port);
		if (localPort == -1) {
			// Use random port
			localPort = 0;
		}
		try {
//			socketTest.bind(new InetSocketAddress(localAddress, localPort));
			//socketTest.connect(stunServer, port);
			while (true) {
				try {
					//traceLog(localAddress, "Remote Address Not Changed: " + socketTest.getRemoteSocketAddress());
					sock.setSoTimeout(timeout);
					//socketTest.connect(stunServer, port);
					MessageHeader sendMH = new MessageHeader(MessageHeader.MessageHeaderType.BindingRequest);
					sendMH.generateTransactionID();
					ChangeRequest changeRequest = new ChangeRequest();
					if (changeAddress) {
						changeRequest.setChangeIP();
						//remoteAddressTemp = changedAddress;
					}
					if (changePort) {
						changeRequest.setChangePort();
						//remotePortTemp = changedPort;
					}
					sendMH.addMessageAttribute(changeRequest);
					byte[] data = sendMH.getBytes();
					DatagramPacket send = new DatagramPacket(data, data.length);
					// socketTest.bind(new InetSocketAddress(stunServer, port));
					send.setAddress(stunServer);
					send.setPort(port);
					sock.send(send);
					traceLog(localAddress, "Test Behave: Binding Request sent.");
					//InetAddress localAddressReal = socketTest.getLocalAddress();
					int localPortReal = sock.getLocalPort();
					traceLog(localAddress, "Local Address not Changed: " + sock.getLocalSocketAddress());
					if (changePort || changeAddress) {
						// socketTest.close();
						// socketTest = new DatagramSocket(new InetSocketAddress(
						// localAddressReal, localPortReal));
						/*
						socketTest.disconnect();
						// socketTest.bind(new InetSocketAddress(remoteAddressTemp, remotePortTemp));
						socketTest.connect(remoteAddressTemp, remotePortTemp);
						traceLog(localAddress, "Remote Address Changed: " + socketTest.getRemoteSocketAddress());
						socketTest.setSoTimeout(timeout);
						socketTest.setReuseAddress(true);
						localAddressReal = socketTest.getLocalAddress();
						localPortReal = socketTest.getLocalPort();
						traceLog(localAddress, "Local Address for Changed: " + socketTest.getLocalSocketAddress());
						*/
					}
					DatagramPacket receive = null;
					MessageHeader receiveMH = new MessageHeader();
					while (!(receiveMH.equalTransactionID(sendMH))) {
						receive = new DatagramPacket(new byte[200], 200);
						sock.receive(receive);
						receiveMH = MessageHeader.parseHeader(receive.getData());
						receiveMH.parseAttributes(receive.getData());
					}
					//localAddressReal = socketTest.getLocalAddress();
					localPortReal = sock.getLocalPort();
					ErrorCode ec = (ErrorCode) receiveMH.getMessageAttribute(MessageAttribute.MessageAttributeType.ErrorCode);
					if (ec != null) {
						di.setError(ec.getResponseCode(), ec.getReason());
						traceLog(localAddress, "Test Behave: Message header contains an Errorcode message attribute.");
						return null;
					} else {
						ChangedAddress ca = (ChangedAddress) receiveMH.getMessageAttribute(MessageAttribute.MessageAttributeType.ChangedAddress);
						MappedAddress ma = (MappedAddress) receiveMH.getMessageAttribute(MessageAttribute.MessageAttributeType.MappedAddress);
						//traceLog(localAddress, "Local Addrress: " + localAddressReal + ":" + localPortReal);
						traceLog(localAddress, "Remote Addrress: " + receive.getAddress() + ":" + receive.getPort());
						traceLog(localAddress, "Changed Address: " + ca);
						traceLog(localAddress, "Mapped Address: " + ma);
						return new ChangeResponse(ca, sock.getInetAddress(), localPortReal, ma);
					}
				} catch (SocketTimeoutException ste) {
					if (timeSinceFirstTransmission < TIME_SINCE_FIRST_TRANSMISSION) {
						traceLog(localAddress, "Test Behave: Socket timeout while receiving the response.");
						timeSinceFirstTransmission += timeout;
						int timeoutAddValue = (timeSinceFirstTransmission * 2);
						if (timeoutAddValue > MAX_TIMEOUT)
							timeoutAddValue = MAX_TIMEOUT;
						timeout = timeoutAddValue;
					} else {
						//socketTest.disconnect();
						//socketTest.close();
						traceLog(localAddress, "Test Behave: Socket timeout while receiving the response. Maximum retry limit exceed. Give up.");
						return new ChangeResponse(true);
					}
				}finally{
					//socketTest.disconnect();
				}
			}
		} finally {
		}
	}

	private void traceLog(String str) {
		logger.debug("Interface:" + iaddress.getHostAddress() + ":" + str);
	}

	private static void traceLog(InetAddress interfaceAddress, String str) {
		logger.debug("Interface:" + interfaceAddress + ":" + str);
	}
}
